import os

import paddlehub as hub
from PIL import Image
import numpy as np


def blend_images(fore_image_path, base_image, out_img_name):
    """
    将抠出的人物图像换背景
    fore_image: 前景图片，抠出的人物图片
    base_image: 背景图片
    """
    # 读入图片
    base_image = Image.open(base_image).convert('RGB')
    fore_image = Image.open(fore_image_path).resize(base_image.size)

    # 图片加权合成
    scope_map = np.array(fore_image)[:, :, -1] / 255
    scope_map = scope_map[:, :, np.newaxis]
    scope_map = np.repeat(scope_map, repeats=3, axis=2)
    res_image = np.multiply(scope_map, np.array(fore_image)[:, :, :3]) + np.multiply((1 - scope_map),
                                                                                     np.array(base_image))
    # 保存图片
    res_image = Image.fromarray(np.uint8(res_image))
    save_path = f'./generated_imgs/{out_img_name}.jpg';
    res_image.save(save_path)
    os.remove(fore_image_path)
    return save_path


def get_head_photo(img_path, img_name, backgroud='white'):
    test_img_path = ["./" + img_path]

    module = hub.Module(name="deeplabv3p_xception65_humanseg", version='1.1.2')

    input_dict = {"image": test_img_path}
    # execute predict and print the result

    module.segmentation(data=input_dict, visualization=True, output_dir="./humanseg_output")
    # blend_images(f'./{img_name}', './backgroud/2.jpg', )
    img_name_text = img_name.split(".")[0]
    if backgroud == 'white':
        back_url = './backgroud/white.png'
    else:
        back_url = './backgroud/2.jpg'
    return blend_images(f'./humanseg_output/{img_name_text}.png', back_url, img_name_text)
